from langchain import hub
from langchain.agents import AgentExecutor, create_react_agent
import os
from dotenv import load_dotenv
from langchain_community.llms import Tongyi

os.environ["DASHSCOPE_API_KEY"] = "sk-9d8f1914800e497f8717144e860f99bc"

model = Tongyi(temperature=1)
from langchain.pydantic_v1 import BaseModel, Field
from langchain.tools import BaseTool, StructuredTool, tool
from langchain.agents import AgentExecutor, create_structured_chat_agent


class CalculatorInput(BaseModel):
    a: str = Field(description="第一个字符串")
    b: str = Field(description="第二个字符串")


def multiply(a: str, b: str) -> int:
    """Multiply two numbers."""
    return len(a) * len(b)


calculator = StructuredTool.from_function(
    func=multiply,  # 工具具体逻辑
    name="Calculator",  # 工具名
    description="计算字符长度的乘积",  # 工具信息
    args_schema=CalculatorInput,  # 工具接受参数信息
    return_direct=True,  # 直接作为工具的输出返回给调用者
    handle_tool_error=True,  # 报错了继续执行，不会吧那些报错行抛出，也可以自定义函数处理，handle_tool_error=函数名
)


class SearchInput(BaseModel):
    query: str = Field(description="should be a search query")


@tool("search-tool", args_schema=SearchInput, return_direct=True)
def search(query: str) -> str:
    """Look up things online."""
    return "你好啊"


tools = [search, calculator]
prompt = hub.pull("hwchase17/structured-chat-agent")
agent = create_structured_chat_agent(model, tools, prompt)
agent_executor = AgentExecutor(
    agent=agent, tools=tools, verbose=True, handle_parsing_errors=True
)
# res = agent_executor.invoke({"input": "`asd`的字符串长度乘以`as`的字符串长度是多少？langchiani是什么？"})
res = agent_executor.invoke({"input": "`asd`的字符串长度乘以`as`的字符串长度是多少？"})
# res = agent_executor.invoke({"input": "langchiani是什么？"})
print(res)